# -*- encoding:utf-8 -*-
"""
-文件：message.py
-说明：消息协议
-作者：tlj
"""
import pickle
import socket
import struct
from block_chain.utxo_set import UTXOSet
from transaction.address import Address
from transaction.transaction import *


# 命令的数据结构，消息体
# version 命令数据类
class CmdVersionBody:
    """握手命令，对比版本和区块链长度，建立链接之后的第一条命令"""
    def __init__(self, source_addr, best_height, version=0):

        self.source_addr = source_addr
        self.best_height = best_height
        self.version = version

    def handler(self, server):
        """
        version消息处理
        :param server:
        :return:
        """
        print(">> 处理 version{}消息".format(self))
        print('='*10)
        #

        bc = server.blockchain
        local_best_height = bc.get_best_height()  # 获取本地节点区块链长度
        peer_best_height = self.best_height  # 对端节点产犊

        if local_best_height<peer_best_height:
            # 获取对端区块链的所有区块的散列值，发送getblocks命令
            BcMsg('getblocks', server.node_address).send_msg(self.source_addr)
        elif local_best_height>peer_best_height:
            # 通过version命令发送本地节点的区块链长度
            BcMsg('version', server.node_address, local_best_height).send_msg(self.source_addr)
        else:
            print("区块链长度一致")
        # 记录对端地址
        if self.source_addr not in server.known_nodes:
            server.known_nodes.append(self.source_addr)


# getdata
class CmdGetDataBody:
    """获取对端节点区块交易数据"""
    def __init__(self, source_addr, data_type, data_id):
        """

        :param source_addr:  发生节点的地址
        :param data_type: 数据类型， 区块block或交易tx
        :param data_id: 如果是区块就是区块的散列值，如果是交易数据就是交易id
        """
        self.source_addr = source_addr
        self.data_type = data_type
        self.data_id = data_id

    def handler(self, server):
        print(">> 处理 getdata{}消息".format(self))
        print('='*10)
        bc = server.blockchain
        if self.data_type == 'block':
            block = bc.get_block(self.data_id)
            if block !=None:
                BcMsg('block', server.node_address, block).send_msg(self.source_addr)
            else:
                print('未找到对应区块！！！')
        elif self.data_type == 'tx':
            try:
                tx = server.mempool[self.data_id]
                BcMsg('tx', server.node_address, tx).send_msg(self.source_addr)
            except KeyError:
                print('未找到对应交易')
        print()


# block
class CmdBlockBody:
    """发送本端节点区块数据命令"""
    def __init__(self, source_addr, block):
        """

        :param source_addr:
        :param block: 区块数据
        """
        self.source_addr = source_addr
        self.block = block

    def handler(self, server):
        print(">> 处理block{}消息".format(self))
        print('='*10)
        print('收到一个新块')
        bc = server.blockchain
        bc.add_block(self.block)
        print('增加一个新块{}'.format(self.block.hash_()))
        if len(server.blocks_in_transit)>0:
            block_hash = server.blocks_in_transit[0]
            BcMsg('getdata', server.node_address, 'block', block_hash).send_msg(self.source_addr)
            server.blocks_in_transit.remove(block_hash)
        else:
            utxo_set = UTXOSet(bc)
            utxo_set.rebuild()
        print()

# inv
class CmdInvBody:
    """发送本段节点的区块散列值或交易id"""
    def __init__(self, source_addr, data_type, items):
        """

        :param source_addr:
        :param data_type: block: 表示发送的是区块散列值，tx：表示发送的是交易id
        :param items:
        """
        self.source_addr = source_addr
        self.data_type = data_type
        self.items = items

    def handler(self, server):
        """
        server.blocks_in_transit 列表保存了本地节点需要下载的所有区块的散列值，
        如果 inv 消息携带的是区块散列值列表，则保存该列表至 server.blocks_in_transit，
        并通过 getdata 消息获取该列表描述的第一个区块。
        我们不一次性获取列表中的所有区块，因为在实际的比特币网络中，
        我们期望能从不同的节点获取区块：
        server.mempool 则保存本地节点的交易池，用于存放其他节点发送的交易，
        如果 inv 消息携带的是交易 ID 列表，且交易 ID 不在交易池中，
        则通过 getdata 消息获取列表中的第一个交易：
        :param server:
        :return:
        """
        print('>> 处理inv 消息{}'.format(self))
        print('='*10)
        if self.data_type == 'block':
            server.blocks_in_transit = self.items
            block_hash = self.items[0]
            BcMsg('getdata', server.node_address, 'block', block_hash).send_msg(self.source_addr)
            server.blocks_in_transit.remove(block_hash)  # 把第一个区块从需要下载的列表中删除
        elif self.data_type == 'tx':
            tx_id = self.items[0]
            if tx_id not in server.mempool.keys():
                BcMsg('getdata', server.node_address, 'tx', tx_id).send_msg(self.source_addr)
                print()


# getblocks
class CmdGetBlocksBody:
    """获取对端节点区块的哈希列表"""
    def __init__(self, source_addr):
        self.source_addr = source_addr

    def handler(self, server):

        print(">> 处理getblocks 消息{}".format(self))
        print('='*10)
        bc = server.blockchain
        blocks_hash = bc.get_blockhashes()
        BcMsg('inv', server.node_address, 'block', blocks_hash).send_msg(self.source_addr)
        print('通过inv命令把本地区块链的所有hash值发送给对端节点')


# tx
class CmdTXBody:
    """发送本段节点交易数据"""
    def __init__(self, source_addr, transaction):
        self.source_addr = source_addr
        self.transaction = transaction

    def handler(self, server):
        print(">> 处理tx{}消息".format(self))
        print("="*10)
        tx = self.transaction
        bc = server.blockchain
        server.mempool[tx.id_] = tx
        print('新增一个交易id:{}到交易池'.format(tx.id_))
        # 如果收到交易信息的是中心节点，则需要向网络中的其他节点发送inv消息告知收到了一个新的交易
        if server.node_address == server.known_nodes[0]:
            for node in server.known_nodes:
                if node!=server.node_address and node!=self.source_addr:
                    BcMsg('inv', server.node_address, 'tx', [tx.id_]).send_msg(node)

        elif len(server.mempool.values()) >1 and len(server.miner_address_str)>0:
            # 如果是矿工节点收到消息, 交易池中交易数量大于1，则打包交易挖出新的区块，
            # 然后通过inv消息，发送给各节点新区块的hash值
            while len(server.mempool.values()) >0:
                txs = []
                tx_ids = []
                for tx_id, tx in server.mempool.items():
                    if bc.verify_transaction(tx):  # 验证交易
                        txs.append(tx)
                        tx_ids.append(tx_id)
                if len(txs) == 0:
                    print('所有交易都无效')
                    return
                cb_tx = CoinbaseTx(Address(address_str=server.miner_address_str))
                txs.append(cb_tx)
                new_block = bc.mine_block(txs)
                utxo_set = UTXOSet(bc)
                utxo_set.update(new_block)
                print("挖出一个新块{}".format(new_block.hash_()))
                for tx_id in tx_ids:
                    del(server.mempool[tx_id])

                for node in server.known_nodes:
                    if node != server.node_address:
                        BcMsg('inv', server.node_address, 'block',
                              [new_block.hash_()]).send_msg(node)






# 消息类定义
# 一个消息类包含一个命令名以及一个对应的命令数据类对象

cmd_body_constructor = {
    'version': CmdVersionBody,
    'block': CmdBlockBody,
    'inv': CmdInvBody,
    'getblocks': CmdGetBlocksBody,
    'getdata': CmdGetDataBody,
    'tx': CmdTXBody,
}


class UnknownMsg(Exception):
    pass


class BcMsg:
    # 消息头12字节长的字节串和一个2字节unsigned short
    # 12字节的命令字符串和两个字节的数据长度组成消息头
    MSG_PACKER = struct.Struct('12p H')
    MSG_HEAD_SIZE = MSG_PACKER.size  # 14字节

    def __init__(self, cmd, *args):
        """
        :param cmd: 命令字符串
        :param args: 消息数据
        """
        self.cmd = cmd
        if self.cmd:
            try:
                constructor = cmd_body_constructor[self.cmd]
                self.body = constructor(*args)
            except KeyError:
                print('未知命令{}'.format(self.cmd))
                raise UnknownMsg
        else:
            self.body = None

    def _pack_msg(self):
        """消息打包"""
        body_bytes = pickle.dumps(self.body)
        # 12字节命令字节串2字节长度然后是数据
        msg = BcMsg.MSG_PACKER.pack(self.cmd.encode('ascii'), len(body_bytes)) + body_bytes
        return msg

    def send_msg(self,target):
        """
        发送数据
        :param target: 目标地址
        :return:
        """
        print("Send msg{} to {}".format(self.cmd, target))
        BcMsg._send_data(BcMsg._unpack_address(target), self._pack_msg())

    @staticmethod
    def _send_data(target, data):
        """发送数据"""
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            sock.connect(target)
            sock.sendall(data)

    @staticmethod
    def _unpack_address(address_str):
        """地址解析"""
        host, port = address_str.split(':')
        port = int(port)
        return (host, port)

    @staticmethod
    def unpack_msg_head(msg_head_bytes):
        """消息头解析"""
        cmd , data_len = BcMsg.MSG_PACKER.unpack(msg_head_bytes)
        cmd = cmd.decode('ascii')
        return cmd, data_len

    def recv_msg(self, the_socket):
        """接收数据"""
        msg_head_bytes = the_socket.recv(BcMsg.MSG_HEAD_SIZE)
        cmd, data_len = BcMsg.unpack_msg_head(msg_head_bytes)
        cmd_body_bytes = the_socket.recv(data_len)
        self.cmd = cmd
        self.body = pickle.loads(cmd_body_bytes)
        return self

    def handler(self, server):
        if hasattr(self.body, 'handler'):
            self.body.handler(server)
        else:
            print('can not handle msg: {}'.format(self.cmd))








